# 数据集类
from typing import Dict
import pickle
import akshare as ak
import numpy as np
import pandas as pd
from pathlib import Path
import subprocess
import shutil
import qlib
from qlib.constant import REG_CN as RegCN
from qlib.data import D

class AksDs(object):
    def __init__(self):
        self.name = 'apps.aks.aks_ds.AksDs'

    @staticmethod
    def get_daily_k(csv_fn:str = 'work/datas/aks/ak_hfq.csv') -> None:
        CODES = ['000001', '000002', '600036', '600519', '300750', '688111']
        START, END = '20200101', '20241231'
        OUT = Path()
        OUT.parent.mkdir(exist_ok=True)
        def limit_ratio(code: str) -> float:    
            return 0.20 if code.startswith(('300', '688')) else (0.05 if code.startswith('ST') else 0.10)
        records = []
        for sym in CODES:    
            df = ak.stock_zh_a_hist(symbol=sym, period='daily', start_date=START, end_date=END, adjust='hfq')    
            df = df.rename(columns={'日期': 'date', '开盘': 'open', '最高': 'high', '最低': 'low', '收盘': 'close','成交量': 'volume', '成交额': 'amount'})    
            df['symbol'] = sym    
            lim = limit_ratio(sym)    
            df['ret'] = df['close'].pct_change()    # (x_{t+1} - x_{t}) / x_{t}
            # df['ret'] = np.log(df['close'] / df['close'].shift(1)) # 对数变化率：log(x_{t+1}/x_{t})
            df['LimitFlag'] = 0    
            df.loc[df['ret'] >= lim - 1e-6, 'LimitFlag'] = 1    
            df.loc[df['ret'] <= -lim + 1e-6, 'LimitFlag'] = -1    
            records.append(df[['date', 'symbol', 'open', 'high', 'low', 'close', 'volume', 'amount', 'LimitFlag']])
        pd.concat(records).to_csv(OUT, index=False)
        print(f'AkShare → CSV 完成：{OUT}  {len(records)} 支股票')

    @staticmethod
    def convert_csv_2_bin(csv_fn:str, bin_folder:str) -> None:
        csv_file = Path(csv_fn)    
        bin_dir  = Path(bin_folder)    
        bin_dir.mkdir(exist_ok=True)    # 1. 临时 csv 目录（dump_bin 要求）
        tmp_csv_dir = bin_dir / 'csv'    
        shutil.rmtree(tmp_csv_dir, ignore_errors=True)    
        tmp_csv_dir.mkdir(parents=True)    
        (tmp_csv_dir / 'ak.csv').write_text(csv_file.read_text())    
        # 2. 调用官方 dump_bin（dump_all 子命令）    
        cmd = [        
            'python', 'qlib/scripts/dump_bin.py', 'dump_all',        
            '--data_path', str(tmp_csv_dir),        
            '--qlib_dir',  str(bin_dir),        
            '--date_field_name', 'date',        
            '--symbol_field_name', 'symbol',        
            '--freq', 'day'    
            ]    
        subprocess.run(cmd, check=True)    
        # 3. 用原始 CSV 生成三列 instrument 文件（防止 qlib 解析错误）    
        qlib.init(provider_uri=str(bin_dir), region=RegCN)    
        df = pd.read_csv(csv_file)    
        df['date'] = pd.to_datetime(df['date'])    
        start_d = df['date'].min().strftime('%Y-%m-%d')    
        end_d   = df['date'].max().strftime('%Y-%m-%d')    
        real_symbols = sorted(df['symbol'].astype(str).unique())    
        split_dir = bin_dir / 'instruments'    
        split_dir.mkdir(exist_ok=True)    
        for split in ['train', 'val', 'all']:        
            (split_dir / f'{split}.txt').write_text(            '\n'.join([f'{sym}\t{start_d}\t{end_d}' for sym in real_symbols]) + '\n'        )    
        print('Three-column instruments written:', split_dir, f'({len(real_symbols)} symbols)')    
        # 4. 清理临时 csv    
        shutil.rmtree(tmp_csv_dir, ignore_errors=True)    
        print('CSV → qlib bin 完成：', bin_dir)

    @staticmethod
    def convert_bin_2_pkl(bin_folder:str, pkl_folder:str) -> None:
        qlib.init(provider_uri=bin_folder, region=RegCN)   # 指向前一步生成的 bin 目录
        # ========================
        out_pkl = Path(pkl_folder)
        out_pkl.mkdir(exist_ok=True)
        def dump(split: str):
            symbols = D.list_instruments({"market": split, "filter_pipe": []})   # ← 传 dict    
            data = {}    
            for sym in symbols:
                df = D.features([sym], ['$open', '$high', '$low', '$close', '$volume', '$amount', '$LimitFlag'],
                                            start_time='2020-01-01', end_time='2024-12-31', freq='day')
                data[sym] = df.dropna()
            with open(out_pkl / f'{split}_data.pkl', 'wb') as f:
                pickle.dump(data, f)    
            print(f'{split} pkl 完成：{len(data)} 符号')
        dump('train')
        dump('val')